import argparse
import re

from build_ases_dml import getPreferredASes

class AttackedAS(object):
    def __init__(self, asn, attacker, victim):
        self.asn = int(asn)
        self.attacker = int(attacker)
        self.victim = int(victim)

        self._was_victim_removed = False
        self._was_attacker_added = False
        self._path_to_victim = []
        self.path_from_attacker_to_victim = []
        self.matched = False

    def set_fwd_table_remove(self, AS_path):
        # Checks that the victim is the last AS in the AS path
        if re.search('{0}\s*$'.format(self.victim), AS_path) is not None:
            self._was_victim_removed = True
        else:
            self._was_victim_removed = False

        self._path_to_victim = []

    def set_fwd_table_add(self, AS_path):
        if re.search('{0}\s*$'.format(self.attacker), AS_path) is not None:
            self._was_attacker_added = True
        else:
            self._was_attacker_added = False

        self._path_to_victim = AS_path.split()

    def get_path_to_victim(self):
        return self._path_to_victim

    def was_affected_by_attack(self):
        return self._was_victim_removed and self._was_attacker_added

    def __eq__(self,other):
        if type(other) == type(self):
            return self.asn == other.asn and self.name == other.name
        if type(other) == int:
            return self.asn == other
        return False
    def __hash__(self):
        return self.asn.__hash__()

bgp_add_regex = lambda as_from, as_to: re.compile('\d*\.\d*\s*bgp@{0}:\d\s*adding.*asp=((?:\d+\s+)*{1}) to'.format(as_from, as_to))

def show_forwarding_table(SSFNET_output_file, asn, *args, **kwargs):
    forwarding_table = get_forwarding_table(SSFNET_output_file, asn)
    print(forwarding_table)

def get_forwarding_table(SSFNET_output_file, as_node):
    '''
    @return a dict, whose keys are ASNs and values are the chosen AS paths.
    '''
    as_node = int(as_node)
    forwarding_table = {}

    with open(SSFNET_output_file, 'r') as ssfnet_file:
        # Find the last path from as_node to any other AS
        for line in ssfnet_file:
            as_node_to_group_regex = bgp_add_regex(as_node, '(\d+)')
            if as_node_to_group_regex.match(line) is not None:
                match = as_node_to_group_regex.match(line)
                AS_path  = match.group(1).split()
                dest = int(match.group(2))
                # Always saving the last one
                forwarding_table[dest] = AS_path

    return forwarding_table

def analyze_attack(SSFNET_output_file, ranking_file, attacker, victim, *args, **kwargs):

    attacker = int(attacker)
    victim = int(victim)
    attacked_ases = []

    with open(SSFNET_output_file, 'r') as ssfnet_file:

        # Get to the attack time
        for line in ssfnet_file:
            # Find the last path from attacker to victim (last is final)
            attacker_to_victim = bgp_add_regex(attacker, victim)
            if attacker_to_victim.match(line) is not None:
                match = attacker_to_victim.match(line)
                attacker_to_victim_AS_path = match.group(1)

            if re.match('Added \d+ to fooled peers', line) is not None:
                break

        for line in ssfnet_file:
            if line[0] in ['|','-']:
                break

            fwd_table_action = re.compile('\d*\.\d* bgp@(\d*):\d\s*(removing|adding).*asp=((?:\d+\s*)+)')
            asn, action, as_path = fwd_table_action.findall(line)[0]
            asn = int(asn)
            if asn in attacked_ases:
                attacked_as = attacked_ases[attacked_ases.index(asn)]
            else:
                attacked_as = AttackedAS(asn, attacker, victim)
                attacked_ases.append(attacked_as)

            if action == 'removing':
                attacked_as.set_fwd_table_remove(as_path)
            else:
                attacked_as.set_fwd_table_add(as_path)

    attack_succeeded = not attacker in attacked_ases

    preferred_ASes = []
    matches = []

    if ranking_file is not None:
        preferred_ASes = getPreferredASes(ranking_file)

    for AS in attacked_ases:
        if attack_succeeded:
            AS.path_from_attacker_to_victim = attacker_to_victim_AS_path.split()
        AS.matched = AS.asn in preferred_ASes
        if AS.matched:
            matches.append(preferred_ASes.index(AS.asn))
        print('{0} affected: {1} path concatation: {2} matched: {3}'.format(AS.asn, AS.was_affected_by_attack(),
                                                                            AS.get_path_to_victim() + AS.path_from_attacker_to_victim, AS.matched))


    print('\r\nAttack succeeded: {0}'.format(attack_succeeded))
    print('\r\nTotal matches: {0} on indices {1}'.format(len(matches), matches))
    return attacked_ases

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("SSFNET_output_file", help="SSFNET execution textual output file")

    subparsers = parser.add_subparsers()

    # Analyze attack subparser
    incoming_parser = subparsers.add_parser('attack', help='Analyzes an attack')
    incoming_parser.set_defaults(func=analyze_attack)
    incoming_parser.add_argument("-r", "--ranking_file", nargs='?', type=argparse.FileType('r'), help="ASes ranking file to check for matches")
    incoming_parser.add_argument("attacker", help="The attacker AS")
    incoming_parser.add_argument("victim", help="The victim AS")

    # Forwarding table subparser
    incoming_parser = subparsers.add_parser('forwarding_table', help='Analyzes an attack')
    incoming_parser.set_defaults(func=show_forwarding_table)
    incoming_parser.add_argument("asn", help="Show this AS forwarding table")


    args = parser.parse_args()
    args.func(**vars(args))